#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/1/23 15:18
# @USER    : Shengji He
# @File    : network_hsj.py
# @Software: PyCharm
# @Version  : Python-
# @TASK:
from pycore.tikzeng import *
from pycore.blocks  import *


def test():
    # arch = [
    #     to_head('..'),      # import package
    #     to_cor(),           # definite color
    #     to_begin(),
    #     to_Conv('conv1', 512, 64, offset="(0, 0, 0)", to="(0, 0, 0)", height=64, depth=64, width=2,),
    #     to_Pool('pool1', offset='(0, 0, 0)', to="(conv1-east)"),
    #     to_Conv("conv2", 128, 64, offset="(1,0,0)", to="(pool1-east)", height=32, depth=32, width=2),
    #     to_connection("pool1", "conv2"),
    #     to_Pool("pool2", offset="(0,0,0)", to="(conv2-east)", height=28, depth=28, width=1),
    #     to_SoftMax("soft1", 10, "(3,0,0)", "(pool1-east)", caption="SOFT"),
    #     to_connection("pool2", "soft1"),
    #     to_Sum("sum1", offset="(1.5,0,0)", to="(soft1-east)", radius=2.5, opacity=0.6),
    #     to_connection("soft1", "sum1"),
    #     to_end()
    # ]
    box_size = 40
    kernel_size = 5      # 4*4
    arch = [
        to_head('..'),
        to_cor(),
        to_begin(),

        # input
        to_input('../examples/fcn8s/cats.jpg'),

        to_ConvRelu(name='cr1', s_filer=28, n_filer=64, offset="(0,0,0)", to="(0,0,0)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr1-east)", height=kernel_size, depth=kernel_size, opacity=0.8),

        to_ConvRelu(name='cr21', s_filer=28, n_filer=64, offset="(5,2,0)", to="(cr1-northeast)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr21-east)", height=kernel_size, depth=kernel_size, opacity=0.8),

        to_ConvRelu(name='cr22', s_filer=28, n_filer=64, offset="(5,-2,0)", to="(cr1-southeast)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr22-east)", height=kernel_size * 2 - 1, depth=kernel_size * 2 - 1, opacity=0.8),

        to_connection("cr1", "cr21"),
        to_connection("cr1", "cr22"),

        to_ConvRelu(name='cr31', s_filer=28, n_filer=64, offset="(3,0,0)", to="(cr21-east)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr31-east)", height=kernel_size, depth=kernel_size, opacity=0.8),

        to_ConvRelu(name='cr32', s_filer=28, n_filer=64, offset="(3,0,0)", to="(cr22-east)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr32-east)", height=kernel_size * 4 - 3, depth=kernel_size * 4 - 3, opacity=0.8),

        to_connection("cr21", "cr31"),
        to_connection("cr22", "cr32"),

        to_ConvRelu(name='cr41', s_filer=28, n_filer=64, offset="(3,0,0)", to="(cr31-east)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr41-east)", height=kernel_size, depth=kernel_size, opacity=0.8),

        to_ConvRelu(name='cr42', s_filer=28, n_filer=64, offset="(3,0,0)", to="(cr32-east)", width=2, height=box_size,
                    depth=box_size),
        to_ConvKernel(name='temp', to="(cr42-east)", height=kernel_size * 8 - 7, depth=kernel_size * 8 - 7, opacity=0.8),

        to_connection("cr31", "cr41"),
        to_connection("cr32", "cr42"),

        to_Cat(name='cr5', s_filer=28, n_filer=128, offset="(17,0,0)", to="(cr1-east)", width=4, height=box_size,
                    depth=box_size),

        to_connection("cr41", "cr5"),
        to_connection("cr42", "cr5"),

        to_Conv(name='cr6', s_filer=28, n_filer=512, offset="(3,0,0)", to="(cr5-east)", width=8, height=box_size,
                    depth=box_size),
        to_Conv(name='cr7', s_filer=28, n_filer=256, offset="(3,0,0)", to="(cr6-east)", width=4, height=box_size,
                depth=box_size),
        to_Conv(name='cr8', s_filer=28, n_filer=3, offset="(3,0,0)", to="(cr7-east)", width=1, height=box_size,
                depth=box_size),

        to_connection("cr5", "cr6"),
        to_connection("cr6", "cr7"),
        to_connection("cr7", "cr8"),
        to_end()
    ]
    to_generate(arch, 'test.tex')


if __name__ == '__main__':
    test()
    print('done')
